//===- basicabc.cpp - Implement basic Array Bounds Check at static time  ---===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//
//
// This file implements two versions of the LLVM "Hello World" pass described
// in docs/WritingAnLLVMPass.html
//
//===----------------------------------------------------------------------===//

#include "llvm/Pass.h"
#include "llvm/Function.h"
#include "llvm/Argument.h"
#include "llvm/Module.h"
#include "llvm/Constants.h"
#include "llvm/Instructions.h"
#include "llvm/LLVMContext.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/Support/InstIterator.h"
#include "llvm/Support/IRBuilder.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Analysis/LoopInfo.h"
#include "llvm/Transforms/Utils/BasicBlockUtils.h"
#include "llvm/Analysis/DebugInfo.h"

#include <iostream>
using namespace llvm;

namespace {

  struct basicabc : public ModulePass {

    static char ID; // Pass identification, replacement for typeid
    IRBuilder<> *builder;
    basicabc() : ModulePass(ID) {}

    virtual bool runOnModule(Module &M) {

      LLVMContext &C = M.getContext();
      Constant *fcall = M.getOrInsertFunction("checkBound",
          /* return type  */  Type::getVoidTy(C), NULL);

      if (!fcall)
      {
        errs() << "Insertion of comparision function failed!\n";
      }

      // set calling convention to C
      Function *checkBound_func = cast<Function>(fcall);
      checkBound_func->setCallingConv(CallingConv::C);
      //errs() << "Argument number " << checkBound_func->arg_size() << "\n";

      BasicBlock* entryBlock = BasicBlock::Create(C, "abcd_entry", checkBound_func);
      builder = new IRBuilder<>(entryBlock);

      //Value* condInst = builder->CreateICmpSLT(checkBound_func->getArgumentList().begin(), checkBound_func->getArgumentList().end(), "cmp1");

      Constant *msg_0 =
        ConstantArray::get(C, "RunTime Warning: Index out of bounds.", true);

      GlobalVariable *aberrormsg = new GlobalVariable(
          M,
          msg_0->getType(),
          true,
          GlobalValue::InternalLinkage,
          msg_0,
          "aberrormsg");

      Function *puts_func = cast<Function>(M.
          getOrInsertFunction("puts", IntegerType::getInt32Ty(C),
            PointerType::getUnqual(IntegerType::getInt8Ty(C)), NULL));

      //errs() << "done insertion\n";
      // error:
      //call i32 @puts(i8 *getelementptr([%d x i8] *@aberrormsg, i32 0, i32 0))
      {
        Constant *zero_32 = Constant::getNullValue(IntegerType::getInt32Ty(C));

        Constant *gep_params[] = {
          zero_32,
          zero_32
        };

        Constant *msgptr = ConstantExpr::
          getGetElementPtr(aberrormsg, gep_params,
              array_lengthof(gep_params));

        Value *puts_params[] = {
          msgptr
        };

        CallInst *puts_call =
          builder->CreateCall(puts_func,
              puts_params, array_endof(puts_params), "put_ret"
              );
        puts_call->setTailCall(false);
      }

      Function *exit_func = cast<Function>(M.
					   getOrInsertFunction("exit", IntegerType::getVoidTy(C), Type::getInt32Ty(C),NULL));

      Value *exit_val = ConstantInt::get(IntegerType::getInt32Ty(C), 1);
      BasicBlock* exitBlock = BasicBlock::Create(C, "abcd_exit", checkBound_func);
      builder->CreateBr(exitBlock);
      delete builder;
      builder = new IRBuilder<>(exitBlock);

      CallInst *exit_call = builder->CreateCall(exit_func,exit_val, "");
      builder->CreateBr(exitBlock);
      delete builder;

      int totalNumberOfChecksInserted = 0, totalNumberOfGEPFound = 0, totalNumberOfArrayTypesFound = 0, totalNumberOfConstIndicesAccessFound = 0;

      for (Module::iterator MI = M.begin(), ME = M.end(); MI != ME; ++MI)
      {
        // if just a function declaration, then continue
        if (MI->isDeclaration())
          continue;

        for (inst_iterator I = inst_begin(*MI), E = inst_end(*MI); I != E; ++I)
        {
          Instruction *inst = &*I;
	  if (isa<UnreachableInst>(inst))
		continue;
          if (isa<GetElementPtrInst>(inst))
          {
            totalNumberOfGEPFound++;
            GetElementPtrInst *GEP = dyn_cast<GetElementPtrInst>(inst);
            if (const ArrayType *ar =
                dyn_cast<ArrayType>(GEP->getPointerOperandType()->getElementType()))
            {
              totalNumberOfArrayTypesFound++;
              unsigned max_elements = ar->getNumElements();

              //errs() << max_elements << " at instrn " << *I << "\n";
              GEPOperator *GEOP = dyn_cast<GEPOperator>(GEP);

              // if all the indices are constant, perform the array bound check statically
              if (GEOP->hasAllConstantIndices())
              {
                totalNumberOfConstIndicesAccessFound++;
                for (unsigned i = 1, e = GEP->getNumOperands(); i != e; ++i) 
                {
                  ConstantInt *v = dyn_cast<ConstantInt>(GEP->getOperand(i));
                  if (!v->getValue().ult(max_elements))
                  {
                    errs() << "CompileTime Warning: Index out of bounds! index " << v->getValue() << " max: " << max_elements << " at instrn " << *I << "\n";

                    if (MDNode *N = I->getMetadata("dbg")) {  // Here I is an LLVM instruction
                      DILocation Loc(N);                      // DILocation is in DebugInfo.h
                      unsigned Line = Loc.getLineNumber();
                      StringRef File = Loc.getFilename();
                      StringRef Dir = Loc.getDirectory();

                      errs() << "At " << Dir << File << ":" << Line << "\n";
                    }
                  }
                }
              }
              // if the indices are non-constant insert code and check
              else
              {
                // insert call to the array bound checking function after this instruction
                // builder.SetInsertPoint(BB);
                
                totalNumberOfChecksInserted++;
                Instruction* next = inst->getNextNode();
                BasicBlock* oldBlock = inst->getParent();
                BasicBlock* newBlock = SplitBlock(oldBlock, next, this);

                int index = GEP->getNumOperands() - 1;
                Value *v1 = GEP->getOperand(index);
                Instruction *i1 = (Instruction *)v1;
                if (isa<SExtInst>(i1) || isa<ZExtInst>(i1))
                {
                  v1 = i1->getOperand(0);
                }
                Value *v2 = ConstantInt::get(v1->getType(), max_elements);

                BasicBlock* checkBlock = BasicBlock::Create(C, "checkBlock", MI, newBlock);
                builder = new IRBuilder<>(oldBlock);
                TerminatorInst* term = oldBlock->getTerminator();
                term->eraseFromParent();
                builder->CreateBr(checkBlock);
                delete builder;
                builder = new IRBuilder<>(checkBlock);
                Value* condInst = builder->CreateICmpSLT(v1, v2, "boundsCheck");
                BasicBlock* errorBlock = BasicBlock::Create(C, "errorBlock", MI, newBlock);
                builder->CreateCondBr(condInst, newBlock, errorBlock);

                delete builder;
                builder = new IRBuilder<>(errorBlock);
                builder->CreateCall(checkBound_func);
                builder->CreateBr(errorBlock);
                delete builder;
              }
            }
          }
        }
      }
      errs() << "The total number of checks inserted are " << totalNumberOfChecksInserted << "\n";
      //errs() << "The total number of pointer accesses are " << totalNumberOfGEPFound << "\n";
      //errs() << "The total number of pointer accesses to arrays are " << totalNumberOfArrayTypesFound << "\n";
      //errs() << "The total number of array accesses where the index is constant  are " << totalNumberOfConstIndicesAccessFound << "\n";
      return true;
    }

    virtual void getAnalysisUsage(AnalysisUsage &AU) const {
      AU.addRequired<LoopInfo>();
    }
  };
}

static void printInst(Instruction* inst)
{
	errs() << *inst << "\n";
	return;
}
char basicabc::ID = 0;
static RegisterPass<basicabc> ABC("basicabc", "basic array bounds check");

